from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
import datetime, os
from typing import cast
from config_types import RootCAConfig


def generate_root_ca(config: RootCAConfig, root_path: str = ""):
    """
    生成根 CA 证书
    
    Args:
        config: 配置字典，包含以下键：
            - output_dir: 输出目录
            - key_path: 私钥文件路径
            - crt_path: 证书文件路径
            - password: 私钥密码
            - country_name: 国家名称
            - state_or_province_name: 省/州名称
            - locality_name: 城市名称
            - organization_name: 组织名称
            - common_name: 通用名称
            - days: 有效期（天数）
            - key_size: RSA密钥长度（可选，默认4096位）
        root_path: 根路径，会与配置中的路径拼接
    """
    KEY_PATH = os.path.join(root_path, config["key_path"]) if root_path else config["key_path"]
    CRT_PATH = os.path.join(root_path, config["crt_path"]) if root_path else config["crt_path"]
    PASSWORD = config["password"].encode()

    COUNTRY_NAME = config["country_name"]
    STATE_OR_PROVINCE_NAME = config["state_or_province_name"]
    LOCALITY_NAME = config["locality_name"]
    ORGANIZATION_NAME = config["organization_name"]
    COMMON_NAME = config["common_name"]

    DAYS = config["days"]
    KEY_SIZE = config.get("key_size", 4096)  # 默认4096位

    # 确保输出目录存在
    if os.path.dirname(KEY_PATH):
        os.makedirs(os.path.dirname(KEY_PATH), exist_ok=True)
    if os.path.dirname(CRT_PATH):
        os.makedirs(os.path.dirname(CRT_PATH), exist_ok=True)

    key = rsa.generate_private_key(public_exponent=65537, key_size=KEY_SIZE)
    subject = issuer = x509.Name(
        [
            x509.NameAttribute(NameOID.COUNTRY_NAME, COUNTRY_NAME),
            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, STATE_OR_PROVINCE_NAME),
            x509.NameAttribute(NameOID.LOCALITY_NAME, LOCALITY_NAME),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, ORGANIZATION_NAME),
            x509.NameAttribute(NameOID.COMMON_NAME, COMMON_NAME),
        ]
    )
    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.utcnow())
        .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=DAYS))
        # ---- 基础扩展 ----
        .add_extension(x509.BasicConstraints(ca=True, path_length=1), critical=True)
        # ---- 权限声明：允许签发下级证书与 CRL ----
        .add_extension(
            x509.KeyUsage(
                digital_signature=True,  # 允许签名
                content_commitment=False,
                key_encipherment=False,
                data_encipherment=False,
                key_agreement=False,
                key_cert_sign=True,  # 允许签发其他证书
                crl_sign=True,  # 允许签发吊销列表
                encipher_only=False,
                decipher_only=False,
            ),
            critical=True,
        )
        # ---- 标识符扩展 ----
        .add_extension(
            x509.SubjectKeyIdentifier.from_public_key(key.public_key()),
            critical=False,
        )
        .add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(key.public_key()),
            critical=False,
        )
        .sign(key, hashes.SHA256())
    )

    # 保存
    with open(KEY_PATH, "wb") as f:
        f.write(
            key.private_bytes(
                serialization.Encoding.PEM,
                serialization.PrivateFormat.TraditionalOpenSSL,
                serialization.BestAvailableEncryption(PASSWORD),
            )
        )
    with open(CRT_PATH, "wb") as f:
        f.write(cert.public_bytes(serialization.Encoding.PEM))

    print("Root CA 生成成功：", CRT_PATH)
    return KEY_PATH, CRT_PATH


if __name__ == "__main__":
    # 独立运行时从配置文件加载
    import json
    from config_types import ConfigDict
    config_path = os.path.join(os.path.dirname(__file__), "./configs/test.json")
    with open(config_path, "r", encoding="utf-8") as f:
        full_config = cast(ConfigDict, json.load(f))
        if "root_ca" not in full_config:
            raise ValueError("配置文件中缺少 'root_ca' 字段")
        config = full_config["root_ca"]
        root_path = full_config.get("root_path", "")
    generate_root_ca(config, root_path)
